#!/usr/bin/python
# -*- coding: UTF-8 -*-
import base64
import os

import qianfan


class AIPattern:
    _instance = None

    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            # 设置环境变量
            os.environ["QIANFAN_ACCESS_KEY"] = "ALTAKVMW05mav4NilQNwpHDlSl"
            os.environ["QIANFAN_SECRET_KEY"] = "0030dc450f60423caeef2b1d390a9572"
        return cls._instance

    def __init__(self, pattern):
        self.model = None
        self.__initialize_model(pattern)

    def __initialize_model(self, pattern):
        if not self.model:
            self.model = self.__choice_model(pattern)

    def __choice_model(self, pattern):
        call_model_method = '_AIPattern__make_' + pattern.lower() + '_model'
        if call_model_method in AIPattern.__dict__:
            model_method = getattr(self, call_model_method)
            if callable(model_method):
                return model_method()

    @staticmethod
    def __make_reply_model():
        print('reply')
        return qianfan.ChatCompletion()

    @staticmethod
    def __make_draw_model():
        print('draw')
        return qianfan.Text2Image()

    def reply(self, wait_msg: str) -> str:
        try:
            return self.model.do(model='ERNIE-4.0-8K', stop=['。'], messages=[{
                'role': 'user',
                'content': wait_msg
            }])['body']['result']
        except Exception as e:
            raise Exception(f'大模型模式错误, 具体原因: {e}')

    def draw(self, keywords):
        try:
            resp = self.model.do(prompt=keywords, with_decode="base64")["body"]["data"][0]["b64_image"]
            return base64.b64decode(resp)
        except Exception as e:
            raise Exception(f'大模型模式错误, 具体原因: {e}')


if __name__ == '__main__':
    ai1 = AIPattern('reply')
    ai2 = AIPattern('draw')
    print(ai1.draw('cat'))
    print(ai1 is ai2)
